package com.spx.chapter05.transform;

import com.spx.chapter05.pojo.Event;
import com.spx.util.SampleDataUtil;
import org.apache.flink.api.common.functions.Partitioner;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;

/**
 * create by undeRdoG on  2022-05-01  13:25
 * 凡心所向，素履以往，生如逆旅，一苇以航。
 */
public class PartitionTest {


    /**
    *  物理分区算子
    * */
    public static void main(String[] args) throws Exception {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        env.setParallelism(1);

        DataStreamSource<Event> dataSource = env.fromCollection(SampleDataUtil.getSample());

        // 当上游和下游算子的并行度调整之后，flink默认会进行分区调整，默认使用的是rebalance
        //dataSource.shuffle().print("shuffle:").setParallelism(4);

        //dataSource.rebalance().print("rebalace:").setParallelism(4);

        /*
        * rescale 重缩放分区，基于rebalance的基础上，在上游的分区的一个分区内rebalance 到 下游多个分区
        * */

        // 这里我们调整source端，创建一个可以并行的source任务
        env.addSource(new RichParallelSourceFunction<Integer>() {
            @Override
            public void run(SourceContext<Integer> sourceContext) throws Exception {
                int indexOfThisSubtask = getRuntimeContext().getIndexOfThisSubtask();
                for (int i = 1; i <= 8; i++){
                    // 根据运行时上下文，获得每个subTask的编号，将奇数放在第一个subTask中，偶数放在第二个subTask中
                    if (i % 2 == indexOfThisSubtask){
                        sourceContext.collect(i);
                    }
                }
            }

            @Override
            public void cancel() {

            }
        }).setParallelism(2).rescale().print("rescale").setParallelism(4);



        // 广播（将一个数据分发到下游所有的并行子任务中）
        dataSource.broadcast();

        // 全局分区 （把所有的数据，都分发到下游任务的第一个并行子任务中）
        dataSource.global();


        // 自定义分区策略
        env.fromElements(1,2,3,4,5,6,7,8).partitionCustom(new Partitioner<Integer>() {
            @Override
            public int partition(Integer key, int numPartitions) {
                return key % 2;
            }
        }, new KeySelector<Integer, Integer>() {
            @Override
            public Integer getKey(Integer value) throws Exception {
                return value;
            }
        }).print().setParallelism(4);

        env.execute();

    }


}
